'''
修改了数据集分割, 训练、验证比例为8:2
'''
import os
import os.path
import numpy as np
from utils import config

IMG_EXTENSIONS = [
    '.jpg', '.JPG', '.jpeg', '.JPEG',
    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]

def is_image_file(filename):
    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)

def dataloader(filepath):
    left_fold = 'Left_Img_Rectified/'
    right_fold = 'Right_Img_Rectified/'
    disp_fold = 'Disp_Map/'

    image = [img for img in os.listdir(filepath+left_fold)]

    img_num = len(image)
    vali_ranges = np.load('./dataloader/USVInland_imgs_ranges.npy')
    if config.kfold != 0:
        print('Kfold:', config.kfold)
        vali_range = vali_ranges[config.kfold-1]
    else:
        vali_range = [0, 1, 4, 12, 21, 27, 29, 31, 43, 47,
            49, 57, 62, 67, 69, 76, 80, 90, 95, 100,
            110, 113, 119, 123, 125, 130, 135, 138, 146, 148,
            158, 159, 163, 167, 170, 176] # validation set indices (36 images in total)

    train = [image[i] for i in range(img_num) if i not in vali_range]
    val = [image[i] for i in vali_range]

    left_train = [filepath+left_fold+img for img in train]
    right_train = [filepath+right_fold+img for img in train]
    disp_train = [filepath+disp_fold+img.replace('jpg', 'png') for img in train]

    left_val = [filepath+left_fold+img for img in val]
    right_val = [filepath+right_fold+img for img in val]
    disp_val = [filepath+disp_fold+img.replace('jpg', 'png') for img in val]

    return left_train, right_train, disp_train, left_val, right_val, disp_val
